#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：torch-cnn1d
@File    ：cnn1d_uniform.py
@IDE     ：PyCharm 
@Author  ：Xiang Qian (chian)
@Date    ：2021/10/6 下午7:44 
'''
import jittor
from ljp import Trainer, layer_cfgs
from ljp.cell import Conv2d_add, Linear_add, Mish

class CNN1D_UNIFORM(Trainer):

    def __init__(self, layer_cfg='4l008', activator='relu', kernel_size=3, n_output=5, operator='subtract',
                 nectar=None):
        super(CNN1D_UNIFORM, self).__init__()
        assert operator in ['subtract', 'mutiply']
        self.arch = 'uni'
        self.operator = operator  # 'std','adder'
        self.activator = activator
        self.layer_cfg = layer_cfg
        self.kernel_size = kernel_size
        self.n_output = n_output
        self.nectar = nectar if nectar else None
        self.init()
        self.desc = '{}-{}-{}-{}-ks{}'.format(self.arch, self.layer_cfg, self.operator, activator, self.kernel_size)

    def init(self):
        self.Conv2D = Conv2d_add if self.operator == 'subtract' else jittor.nn.Conv2d
        cfg_list = layer_cfgs[self.layer_cfg]
        self.activation = jittor.nn.ReLU()
        # self.activation = Mish()
        self.GMP = jittor.nn.AdaptiveAvgPool2d((1, 1))
        softmax_classifier = Linear_add if self.operator == 'subtract' else jittor.nn.Linear
        num_classifier_in = int(
                cfg_list[len(cfg_list) - 2] * self.nectar[len(self.nectar) - 1] / 10) if self.nectar else cfg_list[
            len(cfg_list) - 2]
        self.classifier = softmax_classifier(num_classifier_in, self.n_output, bias=False)
        self.features = self.make_layers(cfg_list)

    def execute(self, x):
        out = self.features(x)
        out = self.GMP(out)
        out = out.view(out.size(0), -1)
        out = self.classifier(out)
        return out

    def make_layers(self, cfg):
        self.conv_padding = 0
        # self.pool_ks = 5
        # self.p_stride = 2
        if self.kernel_size == 5:
            self.pool_ks = 3
            self.p_stride = 3
        elif self.kernel_size == 3:
            self.pool_ks = 3
            self.p_stride = 1
        elif self.kernel_size == 7:
            self.pool_ks = 1
            self.p_stride = 3
        elif self.kernel_size == 9:
            self.pool_ks = 9
            self.conv_padding = 0
            self.p_stride = 2
        else:
            self.pool_ks = 7
            self.p_stride = 3
        layers = []
        in_channels = 1
        Mlayers = 0
        for x_index, x in enumerate(cfg):
            if x == 'm':
                layers += [jittor.nn.AvgPool2d(kernel_size=(1, self.pool_ks), stride=(1, self.p_stride))]
                Mlayers += 1
            else:
                x = int(x * self.nectar[x_index - Mlayers] / 10) if self.nectar else x
                if x == 0:
                    x = 1
                layers += [
                    self.Conv2D(in_channels, x, kernel_size=(1, self.kernel_size), stride=(1, 1), groups=in_channels,
                                padding=(0, self.conv_padding), bias=False),
                    jittor.nn.BatchNorm2d(x),
                    self.activation]
                in_channels = x
        return jittor.nn.Sequential(*layers)
